/*
Copyright (C) 2011 by the PSVN Research Group, University of Alberta
*/

#ifndef _SWITCHBACK_HPP
#define _SWITCHBACK_HPP

#include <vector>
#include <string>
#include "abst_desc.hpp"
#include "astar_data.hpp"
#include "reservoir.hpp"

template <class AH>
class Switchback
{
public:
    explicit Switchback(const AH& ah, bool use_short_circuit);

    ~Switchback();

    bool Solve(const void *orig_state, std::vector<const void*>& solution);

    /* Returns total nodes expanded per level. */
    void NumExpanded(std::vector<size_t>& expanded) const;

    /* Returns total nodes stored per level. */
    void NumStored(std::vector<size_t>& stored) const;

    /* Returns memory required per level. */
    void Memory(std::vector<size_t>& memory) const;

private:
    class Level
    {
    public:
        AStarData data;
        void* work_target;
        void* work_child;
        uint64_t nodes_expanded;
    };

    const AH& m_ah;
    Level* m_levels;
    Reservoir m_memory;
    bool m_use_short_circuit;

    void* NewState();
    int Epsilon(const int i, const void* state) const;

    template<typename GEN>
    const void* Resume(const int i, const void* target);

    int Heuristic(const int i, const void* state);
    void ReconstructPath(const int i, const void* start, 
                          std::vector<const void*>& solution) const;

    /* Non-copyable */
    Switchback(const Switchback& other);
    void operator=(const Switchback& other);
};

template <class AH>
void* Switchback<AH>::NewState()
{
    return m_memory.Get(m_ah.StateSize());
}

template <class AH>
Switchback<AH>::Switchback(const AH& ah, bool use_short_circuit)
    : m_ah(ah),
      m_levels(new Level[ah.NumLevels()]),
      m_memory(16), // allocate in 16MB chunks
      m_use_short_circuit(use_short_circuit)
{
}

template <class AH>
Switchback<AH>::~Switchback()
{
    delete [] m_levels;
}

template <class AH>
int Switchback<AH>::Epsilon(const int i, const void* state) const
{
    // If on odd level: return cost of cheapest rule.
    if ((i & 2) == 1)
        return m_ah.CostOfCheapestBwdRule(i);
    // On even level: 0 if state is goal, otherwise cheapest rule.
    if (m_ah.IsGoal(i, state))
        return 0;
    return m_ah.CostOfCheapestFwdRule(i);
}

template <class AH>
void Switchback<AH>::ReconstructPath(const int i, const void* start, 
                                     std::vector<const void*>& solution) const
{
    const void* current = start;
    std::vector<const void*> sol;
    while (current != NULL) {
        sol.push_back(current);
        StateData data;
        m_levels[i].data.GetData(m_ah.Hash(current), data);
        current = data.From();
    }
    solution = std::vector<const void*>(sol.rbegin(), sol.rend());
}

template <class AH>
int Switchback<AH>::Heuristic(const int i, const void* state)
{
    const int epsilon = Epsilon(i, state);
    if (i == int(m_ah.NumLevels() - 1))
        return epsilon;
    void* abstracted_state = m_levels[i + 1].work_target;
    m_ah.AbstractToNext(i, state, abstracted_state);
    StateData data;
    int64_t abst_hash = m_ah.Hash(abstracted_state);
    if (m_levels[i + 1].data.GetData(abst_hash, data) && data.Expanded())
        return std::max(epsilon, data.G());
    const void* found_goal  = ((i + 1) % 2 == 0) 
        ? Resume<typename AH::FwdChildGenerator>(i + 1, abstracted_state)
        : Resume<typename AH::BwdChildGenerator>(i + 1, abstracted_state);
    if (found_goal != NULL) {
        if (m_levels[i + 1].data.GetData(abst_hash, data)) {
            assert(m_use_short_circuit || data.Expanded());
            return std::max(epsilon, data.G());
        }
        assert(false); // shouldn't happen!
    }
    return INT_MAX;
}

template <class AH>
template<class GEN>
const void* Switchback<AH>::Resume(const int i, const void* target)
{
    // We can also add the short circuit trick here. 
    // Need quick access to fmin (lowest f on open). 
    // I didn't add it here because it looks like an additional 
    // hash look-up is required to get fmin. Is this true?
    while (!m_levels[i].data.OpenListEmpty()) {
        int64_t current_hash;
        StateData current_data;
        m_levels[i].data.Pop(current_hash, current_data);
        m_levels[i].nodes_expanded++;
        void* child = m_levels[i].work_child;
        GEN generator(m_ah, i, current_data.State());
        while (generator.Generate(child)) {

            const int64_t child_hash = m_ah.Hash(child);
            const int child_g = current_data.G() + generator.Cost();
            // relax previously seen child
            if (m_levels[i].data.Relax(child_hash, child_g, 
                                       current_data.State())) {
                continue;
            }
            // add new child
            void* new_child = NewState();
            m_ah.Copy(new_child, child);
            const int child_h = Heuristic(i, new_child);
            assert(child_h != INT_MAX);
            m_levels[i].data.AddData(child_hash, 
                                     StateData(new_child, 
                                               current_data.State(),
                                               child_g, child_h),
                                     child_g + child_h, child_g);
        }

	// Short Circuit. 
	if(target != NULL && m_use_short_circuit){
	    
            // Search open for the target. 
            int64_t target_hash = m_ah.Hash(target);
            StateData sd;
            if(m_levels[i].data.GetData(target_hash, sd)) {
                // Check its bounds (see paper). 
                if(( (double)sd.G() / (current_data.G() + current_data.H() 
                                       - (double)sd.H())) <= 1)
                {
                    return sd.State();
                }
            }
	}

        if (i == 0) {
            if (m_ah.IsGoal(i, current_data.State()))
                return current_data.State();
        } else if (m_ah.Equal(current_data.State(), target))
            return current_data.State();
    }
    return NULL;
}

template <class AH>
bool Switchback<AH>::Solve(const void* orig_state, 
                           std::vector<const void*>& solution) 
{    
    if (m_ah.IsGoal(0, orig_state))
        return true;
    void* cur_state = NULL;
    m_memory.Reset();
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        m_levels[i].work_target = NewState();
        m_levels[i].work_child = NewState();
        m_levels[i].data.Clear();
        m_levels[i].nodes_expanded = 0;
        // copy over start state if at base level
        if (i == 0) {
            cur_state = NewState();
            m_ah.Copy(cur_state, orig_state);
        }
        // abstract start state to this level
        if (i > 0) { 
            void* prev_state = cur_state;
            cur_state = NewState();
            m_ah.AbstractToNext(i - 1, prev_state, cur_state);
        }
        // add states to open list
        if (i % 2 == 0) {
            int64_t hash = m_ah.Hash(cur_state);
            m_levels[i].data.AddData(hash, StateData(cur_state), 0, 0);
        } else {
            int n = 0;
            void* goal_state_iter = m_levels[i].work_target;
            m_ah.CopyGoalState(i, goal_state_iter, 0);
            do {
                // Need to create a copy so the entry in the 
                // hashtable is pointing to something valid.
                void* goal = NewState();
                m_ah.Copy(goal, goal_state_iter);
                int64_t hash = m_ah.Hash(goal);
                m_levels[i].data.AddData(hash, StateData(goal), 0, 0);
            } while (m_ah.NextGoalState(i, goal_state_iter, &n));
        }
    }
    // I think passing a null here will prevent short circuit from 
    // working at the base level. Could we pass the goal instead? 
    void* target = NewState();
    m_ah.CopyGoalState(0, target, 0);
    const void* goal_state = Resume<typename AH::FwdChildGenerator>(0, target);
    if (goal_state) {
        ReconstructPath(0, goal_state, solution);
        return true;
    }
    return false;
}

template <class AH>
void Switchback<AH>::NumExpanded(std::vector<size_t>& expanded) const
{
    expanded.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        expanded[i] = m_levels[i].nodes_expanded;
}

template <class AH>
void Switchback<AH>::NumStored(std::vector<size_t>& stored) const
{
    stored.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        stored[i] = m_levels[i].data.NumStored();
}

template <class AH>
void Switchback<AH>::Memory(std::vector<size_t>& mem) const
{
    mem.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        mem[i] = m_levels[i].data.Memory() 
            + m_levels[i].data.NumStored() * m_ah.StateSize();
    }
}

#endif // _SWITCHBACK_HPP
